package com.abyss.state;


import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple1;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.datastream.BroadcastConnectedStream;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.util.Collector;

import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * 需求：公司有10个广告位, 其广告的内容(描述和图片)会经常变动（广告到期，更换广告等）
 * 实现：
 * 1. 通过socket流输入广告id
 * 2. 关联出广告的信息打印出来，如果广告发生变化，需要能够感知到
 *    以一个自定义Source来模拟这个过程（随机数）
 *    要注意：真实业务场景中这种变化数据的获取一般通过内网API或者redis等外部系统动态查询
 */
public class BroadcastStateDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.enableCheckpointing(5000L, CheckpointingMode.EXACTLY_ONCE);
        env.setStateBackend(new FsStateBackend("file:///Users/abyss/Dev/toys/flink/H-flink-learn/src/main/resources/br_state_ck"));

        // source
        SingleOutputStreamOperator<Tuple1<Integer>> idData = env.socketTextStream("node1", 9999).map(new MapFunction<String, Tuple1<Integer>>() {
            @Override
            public Tuple1<Integer> map(String value) throws Exception {
                return Tuple1.of(Integer.parseInt(value));
            }
        });

        // 获得广告的source对象
        DataStreamSource<Map<Integer, Tuple2<String, String>>> adSource = env.addSource(new MySourceForBroadcast());

        // 将广告数据广播出去广播状态
        // 得到一个广播流
        BroadcastStream<Map<Integer, Tuple2<String, String>>> broadcastStream = adSource.broadcast(new MapStateDescriptor<Integer, Map<Integer, Tuple2<String, String>>>(
                "broadcast state",
                TypeInformation.of(new TypeHint<Integer>() {
                }),
                TypeInformation.of(new TypeHint<Map<Integer, Tuple2<String, String>>>() {
                })
        ));

        // 让数据流和广播流 关联到一起，让这个广播流一直跟随数据流
        BroadcastConnectedStream<Tuple1<Integer>, Map<Integer, Tuple2<String, String>>> broadcastConnectedStream = idData.connect(broadcastStream);

        // 对得到的BroadcastConnectedStream进行处理
        SingleOutputStreamOperator<Tuple2<String, String>> result = broadcastConnectedStream.process(new BroadcastProcessFunction<Tuple1<Integer>, Map<Integer, Tuple2<String, String>>, Tuple2<String, String>>() {
            private MapStateDescriptor<Integer, Map<Integer, Tuple2<String, String>>> mapMapStateDescriptor = new MapStateDescriptor<>(
                    "broadcast state",
                    TypeInformation.of(new TypeHint<Integer>() {
                    }),
                    TypeInformation.of(new TypeHint<Map<Integer, Tuple2<String, String>>>() {
                    })
            );

            private BroadcastState<Integer, Map<Integer, Tuple2<String, String>>> mapState;

            /*
            这个方法处理的是数据流， 会得到一个只读的context，从context中可以去到mapState（只读不可写）
            取到后，就可以取出里面记录的最新的数据了
            拿到数据自行操作
             */
            @Override
            public void processElement(Tuple1<Integer> id, ReadOnlyContext ctx, Collector<Tuple2<String, String>> out) throws Exception {
                ReadOnlyBroadcastState<Integer, Map<Integer, Tuple2<String, String>>> readOnlyBroadcastState = ctx.getBroadcastState(mapMapStateDescriptor);
                Map<Integer, Tuple2<String, String>> map = readOnlyBroadcastState.get(1);
                Tuple2<String, String> result = map.get(id.f0);
                if (result != null) out.collect(result);
            }

            /*
            将广播流source每隔5秒产生的数据作为参数1传递进来
            通过参数2 context对象取得mapState，将参数1最新的数据 塞入mapState中
            供上面的processElement方法使用
             */
            @Override
            public void processBroadcastElement(Map<Integer, Tuple2<String, String>> value, Context ctx, Collector<Tuple2<String, String>> out) throws Exception {
                mapState = ctx.getBroadcastState(mapMapStateDescriptor);
                mapState.clear();
                mapState.put(1, value);
            }
        });

        result.print();
        env.execute();

    }

    /*
    自定义source
    每隔5秒产生一个map，map里面记录了10条广告数据，key是广告id 从1-10， value是广播信息是一个Tuple2<公司名称, 广告语>
     */
    public static class MySourceForBroadcast implements SourceFunction<Map<Integer, Tuple2<String, String>>> {
        private final Random random = new Random();
        private final List<Tuple2<String, String>> ads = Arrays.asList(
                Tuple2.of("baidu", "搜索引擎"),
                Tuple2.of("google", "科技大牛"),
                Tuple2.of("aws", "全球领先的云平台"),
                Tuple2.of("aliyun", "全球领先的云平台"),
                Tuple2.of("腾讯", "氪金使我变强"),
                Tuple2.of("阿里巴巴", "电商龙头"),
                Tuple2.of("字节跳动", "靠算法出名"),
                Tuple2.of("美团", "黄色小公司"),
                Tuple2.of("饿了么", "蓝色小公司"),
                Tuple2.of("瑞幸咖啡", "就是好喝")
        );
        private boolean isRun = true;

        @Override
        public void run(SourceContext<Map<Integer, Tuple2<String, String>>> ctx) throws Exception {
            while (isRun) {
                Map<Integer, Tuple2<String, String>> map = new HashMap<>();
                int keyCounter = 0;
                for (int i = 0; i < ads.size(); i++) {
                    keyCounter++;
                    map.put(keyCounter, ads.get(random.nextInt(ads.size())));
                }
                ctx.collect(map);

                TimeUnit.SECONDS.sleep(5L);
            }
        }

        @Override
        public void cancel() {
            this.isRun = false;
        }
    }
}